"""
12306官方API集成（无需申请，直接使用）
"""
import requests
import json
from datetime import datetime, timedelta

class RailwayAPI:
    def __init__(self):
        self.base_url = "https://kyfw.12306.cn/otn/leftTicket"
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0 Safari/537.36',
            'Referer': 'https://kyfw.12306.cn/otn/leftTicket/init',
            'Accept': 'application/json, text/javascript, */*; q=0.01',
            'Accept-Language': 'zh-CN,zh;q=0.9',
            'X-Requested-With': 'XMLHttpRequest',
            'Connection': 'keep-alive'
        }
        self.session = requests.Session()
        try:
            # 预热，获取必要的Cookie
            self.session.get('https://kyfw.12306.cn/otn/leftTicket/init', headers=self.headers, timeout=8, verify=False)
        except Exception:
            pass
        # 关闭证书告警
        try:
            requests.packages.urllib3.disable_warnings()  # type: ignore
        except Exception:
            pass
    
    def get_train_info(self, origin, destination, date=None):
        """获取火车信息（增强容错）"""
        if not date:
            date = (datetime.now() + timedelta(days=1)).strftime('%Y-%m-%d')
        
        origin_code = self.get_station_code(origin)
        dest_code = self.get_station_code(destination)
        if not origin_code or not dest_code:
            return self.get_mock_data(origin, destination)
        
        params = {
            'leftTicketDTO.train_date': date,
            'leftTicketDTO.from_station': origin_code,
            'leftTicketDTO.to_station': dest_code,
            'purpose_codes': 'ADULT'
        }
        endpoints = ['queryZ', 'query', 'queryA']
        for ep in endpoints:
            try:
                url = f"https://kyfw.12306.cn/otn/leftTicket/{ep}"
                r = self.session.get(url, params=params, headers=self.headers, timeout=10, verify=False)
                text = r.text
                # 兼容返回体
                try:
                    data = r.json()
                except Exception:
                    # 部分情况下返回非JSON（如被拦），尝试截断调试
                    print(f"获取火车数据失败: 非JSON响应，状态{r.status_code}，片段: {text[:120]}")
                    continue
                if data and data.get('status') and data.get('data'):
                    raw = data['data']
                    results = raw.get('result') or raw  # 有的版本直接在data里
                    return self.parse_train_data(results, origin, destination)
            except Exception as e:
                print(f"获取火车数据失败: {e}")
                continue
        # 全部失败回退
        return self.get_mock_data(origin, destination)
    
    def get_station_code(self, city):
        """获取车站代码"""
        station_codes = {
            '北京': 'BJP',
            '上海': 'SHH', 
            '广州': 'GZQ',
            '深圳': 'SZN',
            '杭州': 'HZH',
            '南京': 'NJH',
            '武汉': 'WHN',
            '成都': 'CDW',
            '西安': 'XAY',
            '重庆': 'CQW',
            '天津': 'TJP',
            '青岛': 'QDK',
            '大连': 'DLT',
            '沈阳': 'SYT',
            '哈尔滨': 'HBB',
            '长春': 'CCT',
            '石家庄': 'SJP',
            '太原': 'TYV',
            '呼和浩特': 'HHC',
            '银川': 'YIJ',
            '西宁': 'XNO',
            '兰州': 'LZJ',
            '乌鲁木齐': 'WAR',
            '拉萨': 'LSO',
            '昆明': 'KMM',
            '贵阳': 'GIW',
            '南宁': 'NNZ',
            '海口': 'VUQ',
            '三亚': 'SEQ',
            '福州': 'FZS',
            '厦门': 'XMS',
            '南昌': 'NCG',
            '长沙': 'CSQ',
            '合肥': 'HFH',
            '郑州': 'ZZF',
            '济南': 'JNK',
            '太原': 'TYV'
        }
        return station_codes.get(city, 'BJP')
    
    def parse_train_data(self, data, origin, destination):
        """解析火车数据，兼容多种返回格式"""
        trains = []
        # 可能是列表result里的竖线分隔字符串
        if isinstance(data, list) and data and isinstance(data[0], str):
            for train_info in data[:10]:
                try:
                    parts = train_info.split('|')
                    if len(parts) < 10:
                        continue
                    trains.append({
                        'type': 'high_speed_rail' if parts[3].startswith(('G','D')) else 'train',
                        'train_number': parts[3],
                        'departure_time': parts[8],
                        'arrival_time': parts[9],
                        'duration': parts[10],
                        'price': 0,
                        'seats_available': parts[30] if len(parts) > 30 else '',
                        'origin': origin,
                        'destination': destination
                    })
                except Exception:
                    continue
            return trains[:5]
        # 也可能已经是解析后的对象列表
        if isinstance(data, list) and data and isinstance(data[0], dict):
            for item in data[:10]:
                q = item.get('queryLeftNewDTO') or item
                code = q.get('station_train_code', '')
                if not code:
                    continue
                trains.append({
                    'type': 'high_speed_rail' if code.startswith(('G','D')) else 'train',
                    'train_number': code,
                    'departure_time': q.get('start_time', ''),
                    'arrival_time': q.get('arrive_time', ''),
                    'duration': q.get('lishi', ''),
                    'price': q.get('zy_price', 0) or q.get('ze_price', 0) or 0,
                    'seats_available': q.get('zy_num', '') or q.get('ze_num', ''),
                    'origin': q.get('from_station_name', origin),
                    'destination': q.get('to_station_name', destination)
                })
            return trains[:5]
        return []
    
    def get_mock_data(self, origin, destination):
        """模拟数据"""
        return [{
            'type': 'high_speed_rail',
            'train_number': 'G1',
            'departure_time': '08:00',
            'arrival_time': '12:30',
            'duration': '4.5小时',
            'price': 553,
            'seats_available': '有票',
            'origin': origin,
            'destination': destination
        }]

# 使用示例
if __name__ == "__main__":
    api = RailwayAPI()
    trains = api.get_train_info('北京', '上海')
    print("火车信息:", json.dumps(trains, ensure_ascii=False, indent=2))